/****************************************************************************
Project: lio-sam-ros2
Github: https://gitee.com/liu-fang66/lio-sam-ros2
Author: liufang
EMail: 543334547@qq.com
Date: 2024-12-12
-------------------------------------------------------
*featureExtraction*

功能：
+ 对经过运动畸变校正之后的当前帧激光点云，计算每个点的曲率，进而提取角点、平面点（用曲率的大小进行判定）

订阅：
1.
去畸变后的点云及附属信息，包括：原始点云、去畸变点云、该帧点云的初始旋转旋转角（来自IMU原始roll、pitch、yaw）、该帧点云的初始位姿（来自IMU里程计）

发布：
1. 包含提取完特征点的点云信息包。在订阅的信息包基础上塞进去提取出的特征角点和特征平面点

流程：
1. 接收到从imageProjection中发布出的一个去畸变点云信息cloudInfo(自定义格式)
2. 对每个点计算曲率。计算时是计算周围点的平均距离用来作为曲率的替代
3. 标记遮挡点和与激光平行的点，后续这些点不能采纳为特征点
4. 特征提取。分别做角点（曲率大）和平面点（曲率小）特征点提取
5. 整合信息，发布完整数据包
*******************************************************************************/

#include "lio_sam/msg/cloud_info.hpp"
#include "utility.hpp"

struct smoothness_t {
    float  value;  // 曲率值
    size_t ind;    // 激光点一维索引
};

// 激光点曲率排序函数，从小到大排序
struct by_value {
    bool operator()(smoothness_t const &left, smoothness_t const &right) { return left.value < right.value; }
};

class FeatureExtraction : public ParamServer {
public:
    rclcpp::Subscription<lio_sam::msg::CloudInfo>::SharedPtr p_sub_laser_cloud_Info_;

    // 发布当前激光帧提取特征之后的点云信息
    rclcpp::Publisher<lio_sam::msg::CloudInfo>::SharedPtr p_pub_laser_cloud_info_;
    // 发布当前激光帧提取的角点点云
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr p_pub_corner_points_;
    // 发布当前激光帧提取的平面点点云
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr p_pub_surface_points_;

    pcl::PointCloud<PointType>::Ptr p_extracted_cloud_;  // 当前激光帧运动畸变校正后的有效点云
    pcl::PointCloud<PointType>::Ptr p_corner_cloud_;     // 当前激光帧角点点云集合
    pcl::PointCloud<PointType>::Ptr p_surface_cloud_;    // 当前激光帧平面点点云集合

    pcl::VoxelGrid<PointType> down_size_filter_;  // 用来做平面特征点点云降采样的过滤器

    // 当前激光帧点云信息，包括的历史数据有：运动畸变校正，点云数据，初始位姿，姿态角，有效点云数据，角点点云，平面点点云等
    lio_sam::msg::CloudInfo cloud_info_;
    std_msgs::msg::Header   cloud_header_;  // 接收到的当前帧信息的数据头

    // 点云曲率，cloudSmoothness可以用来做排序，所以里面有另一个字段存储index
    std::vector<smoothness_t> cloud_smoothness_vec_;
    float                    *p_cloud_curvature_;  // 用来做曲率计算的中间变量

    // 0表示还未进行特征提取处理,1表示遮挡、平行，或者已经进行特征提取的点，一个N_SCAN*Horizon_SCAN长的数组
    int *p_cloud_neighbor_picked_;
    int *p_cloud_label_;  // 1表示角点，-1表示平面点，一个N_SCAN*Horizon_SCAN长的数组

    FeatureExtraction(const rclcpp::NodeOptions &options) : ParamServer("lio_sam_featureExtraction", options) {
        // 订阅当前激光帧运动畸变校正后的点云信息
        p_sub_laser_cloud_Info_ = create_subscription<lio_sam::msg::CloudInfo>(
            "lio_sam/deskew/cloud_info", qos,
            std::bind(&FeatureExtraction::laserCloudInfoHandler, this, std::placeholders::_1));

        // 发布当前激光帧提取特征之后的点云信息
        p_pub_laser_cloud_info_ = create_publisher<lio_sam::msg::CloudInfo>("lio_sam/feature/cloud_info", qos);

        // 订阅当前激光帧运动畸变校正后的点云信息
        p_pub_corner_points_ = create_publisher<sensor_msgs::msg::PointCloud2>("lio_sam/feature/cloud_corner", 1);

        // 发布当前激光帧的面点点云
        p_pub_surface_points_ = create_publisher<sensor_msgs::msg::PointCloud2>("lio_sam/feature/cloud_surface", 1);

        initializationValue();
    }

    void initializationValue() {
        cloud_smoothness_vec_.resize(N_SCAN * Horizon_SCAN);

        down_size_filter_.setLeafSize(odometry_surf_leaf_size_, odometry_surf_leaf_size_, odometry_surf_leaf_size_);

        p_extracted_cloud_.reset(new pcl::PointCloud<PointType>());
        p_corner_cloud_.reset(new pcl::PointCloud<PointType>());
        p_surface_cloud_.reset(new pcl::PointCloud<PointType>());

        p_cloud_curvature_       = new float[N_SCAN * Horizon_SCAN];
        p_cloud_neighbor_picked_ = new int[N_SCAN * Horizon_SCAN];
        p_cloud_label_           = new int[N_SCAN * Horizon_SCAN];
    }

    /**
     * @brief
     * cloudInfo话题的回调函数，这个模块的功能都是顺序进行
     * 1. 接收到从imageProjection中发布出的一个去畸变点云信息cloudInfo(自定义格式)
     * 2. 对每个点计算曲率。计算时是计算周围点的平均距离用来作为曲率的替代
     * 3. 标记遮挡点和与激光平行的点，后续这些点不能采纳为特征点
     * 4. 特征提取。分别做角点（曲率大）和平面点（曲率小）特征点提取
     * 5. 整合信息，发布完整数据包
     *
     * @param msgIn 从去畸变模块接受的数据包
     */
    void laserCloudInfoHandler(const lio_sam::msg::CloudInfo::SharedPtr msgIn) {
        cloud_info_   = *msgIn;         // new cloud info
        cloud_header_ = msgIn->header;  // new cloud header
        pcl::fromROSMsg(msgIn->cloud_deskewed,
                        *p_extracted_cloud_);  // new cloud for extraction

        calculateSmoothness();  // 计算当前激光帧点云中每个点的曲率

        markOccludedPoints();  // 标记属于遮挡、平行两种情况的点，不做特征提取

        // 点云角点、平面点特征提取
        // 1、遍历扫描线，每根扫描线扫描一周的点云划分为6段，针对每段提取20个角点、不限数量的平面点，加入角点集合、平面点集合
        // 2、认为非角点的点都是平面点，加入平面点云集合，最后降采样
        extractFeatures();

        publishFeatureCloud();  // 发布角点、面点点云，发布带特征点云数据的当前激光帧点云信息
    }

    void calculateSmoothness() {
        int cloudSize = p_extracted_cloud_->points.size();

        // 遍历当前激光帧运动畸变校正后的有效点云
        for (int i = 5; i < cloudSize - 5; i++) {
            // 用当前激光点前后5个点计算当前点的曲率
            // 注意，这里把前后五个点共10个点加起来，还减去了10倍的当前点
            float diffRange =
                cloud_info_.point_range[i - 5] + cloud_info_.point_range[i - 4] + cloud_info_.point_range[i - 3] +
                cloud_info_.point_range[i - 2] + cloud_info_.point_range[i - 1] - cloud_info_.point_range[i] * 10 +
                cloud_info_.point_range[i + 1] + cloud_info_.point_range[i + 2] + cloud_info_.point_range[i + 3] +
                cloud_info_.point_range[i + 4] + cloud_info_.point_range[i + 5];

            // 距离差值平方作为曲率
            p_cloud_curvature_[i] = diffRange * diffRange;  // diffX * diffX + diffY * diffY + diffZ * diffZ;

            p_cloud_neighbor_picked_[i] = 0;
            p_cloud_label_[i]           = 0;
            // cloud_smoothness_vec_ for sorting
            // 存储该点曲率值、激光点一维索引
            // 之所以可以这样操作，是因为在initializationValue部分，对cloudSmoothness进行过初始化，
            // 否则直接对cloudSmoothness[i]赋值，一定会报段错误
            cloud_smoothness_vec_[i].value = p_cloud_curvature_[i];
            cloud_smoothness_vec_[i].ind   = i;
        }
    }

    void markOccludedPoints() {
        int cloudSize = p_extracted_cloud_->points.size();
        // mark occluded points and parallel beam points
        for (int i = 5; i < cloudSize - 6; ++i) {
            // occluded points
            // 当前点和下一个点的range值
            float depth1 = cloud_info_.point_range[i];
            float depth2 = cloud_info_.point_range[i + 1];
            // 两个激光点之间的一维索引差值，如果在一条扫描线上，那么值为1；
            // 如果两个点之间有一些无效点被剔除了，可能会比1大，但不会特别大
            // 如果恰好前一个点在扫描一周的结束时刻，下一个点是另一条扫描线的起始时刻，那么值会很大
            int columnDiff = std::abs(int(cloud_info_.point_col_ind[i + 1] - cloud_info_.point_col_ind[i]));

            // 两个点在同一扫描线上，且距离相差大于0.3，认为存在遮挡关系
            // （也就是这两个点不在同一平面上，如果在同一平面上，距离相差不会太大）
            //  远处的点会被遮挡，标记一下该点以及相邻的5个点，后面不再进行特征提取
            if (columnDiff < 10) {
                // 10 pixel diff in range image
                if (depth1 - depth2 > 0.3) {
                    p_cloud_neighbor_picked_[i - 5] = 1;
                    p_cloud_neighbor_picked_[i - 4] = 1;
                    p_cloud_neighbor_picked_[i - 3] = 1;
                    p_cloud_neighbor_picked_[i - 2] = 1;
                    p_cloud_neighbor_picked_[i - 1] = 1;
                    p_cloud_neighbor_picked_[i]     = 1;
                } else if (depth2 - depth1 > 0.3) {
                    p_cloud_neighbor_picked_[i + 1] = 1;
                    p_cloud_neighbor_picked_[i + 2] = 1;
                    p_cloud_neighbor_picked_[i + 3] = 1;
                    p_cloud_neighbor_picked_[i + 4] = 1;
                    p_cloud_neighbor_picked_[i + 5] = 1;
                    p_cloud_neighbor_picked_[i + 6] = 1;
                }
            }
            // parallel beam
            // 用前后相邻点判断当前点所在平面是否与激光束方向平行
            // diff1和diff2是当前点距离前后两个点的距离
            float diff1 = std::abs(float(cloud_info_.point_range[i - 1] - cloud_info_.point_range[i]));
            float diff2 = std::abs(float(cloud_info_.point_range[i + 1] - cloud_info_.point_range[i]));

            // 如果当前点距离左右邻点都过远，则视其为瑕点，因为入射角可能太小导致误差较大
            //  选择距离变化较大的点，并将他们标记为1
            if (diff1 > 0.02 * cloud_info_.point_range[i] && diff2 > 0.02 * cloud_info_.point_range[i])
                p_cloud_neighbor_picked_[i] = 1;
        }
    }

    /**
     * @brief
     * 特征点提取，遍历扫描线，每根扫描线扫描一周的点云划分为6段
     * 1. 提取角点
     *      1.1 针对每段提取最多20个角点
     *      1.2 设置cloudLabel标志，cloudLabel=1标志该点为角点
     *      1.3 cloudNeighborPicked=1标记该点已进行特征提取
     *      1.4 对该角点左右各5个点，如果两点之间的列索引差距小于10，则抛弃周围的点，避免重复对同一块区域提取角点
     *      1.5 将该角点加入角点点云集合
     * 2. 提取平面点
     *      2.1 不限制每段平面点的数量，后面会进行降采样
     *      2.2 设置cloudLabel标志，cloudLabel=-1标志该点为平面点
     *      2.3 cloudNeighborPicked=1标记该点已进行特征提取
     *      2.3 对该平面点左右各5个点，如果两点之间的列索引差距小于10，则抛弃周围的点，避免重复对同一块区域提取角点
     * 3. 将该段中除了角点之外的点加入平面点集合
     *      ！！！ 这点让步骤2感觉是多余的，最终的结果可能只是原始点云降采样，可能特征点提取也是没有必要的！！！
     * 4. 遍历结束
     * 5. 将平面特征点集合进行降采样
     *
     */
    void extractFeatures() {
        p_corner_cloud_->clear();
        p_surface_cloud_->clear();

        pcl::PointCloud<PointType>::Ptr surfaceCloudScan(new pcl::PointCloud<PointType>());
        pcl::PointCloud<PointType>::Ptr surfaceCloudScanDS(new pcl::PointCloud<PointType>());

        for (int i = 0; i < N_SCAN; i++) {
            surfaceCloudScan->clear();

            // 将一条扫描线扫描一周的点云数据，划分为6段，每段分开提取有限数量的特征，保证特征均匀分布
            for (int j = 0; j < 6; j++) {
                // 每段点云的起始、结束索引；startRingIndex为扫描线起始第5个激光点在一维数组中的索引
                // 注意：所有的点云在这里都是以"一维数组"的形式保存
                // startRingIndex和 endRingIndex 在imageProjection.cpp中的cloudExtraction函数里被填入
                // 假设 当前ring在一维数组中起始点是m，结尾点为n（不包括n），那么6段的起始点分别为：
                // m + [(n-m)/6]*j   j从0～5
                // 化简为 [（6-j)*m + nj ]/6
                // 6段的终止点分别为：
                // m + (n-m)/6 + [(n-m)/6]*j -1  j从0～5,-1是因为最后一个,减去1
                // 化简为 [（5-j)*m + (j+1)*n ]/6 -1
                // 这块不必细究边缘值到底是不是划分的准（例如考虑前五个点是不是都不要，还是说只不要前四个点），
                // 只是尽可能的分开成六段，首位相接的地方不要。因为庞大的点云中，一两个点其实无关紧要。
                int sp = (cloud_info_.start_ring_index[i] * (6 - j) + cloud_info_.end_ring_index[i] * j) / 6;
                int ep = (cloud_info_.start_ring_index[i] * (5 - j) + cloud_info_.end_ring_index[i] * (j + 1)) / 6 - 1;

                if (sp >= ep)
                    continue;

                // 按照曲率从小到大排序点云
                // 可以看出之前的byvalue在这里被当成了判断函数来用
                std::sort(cloud_smoothness_vec_.begin() + sp, cloud_smoothness_vec_.begin() + ep, by_value());

                int largestPickedNum = 0;
                // 按照曲率从大到小遍历
                for (int k = ep; k >= sp; k--) {
                    int ind = cloud_smoothness_vec_[k].ind;  // 激光点的索引

                    // 当前激光点还未被处理，且曲率大于阈值，则认为是角点
                    if (p_cloud_neighbor_picked_[ind] == 0 && p_cloud_curvature_[ind] > edge_threshold_) {
                        largestPickedNum++;

                        // 每段只取20个角点，如果单条扫描线扫描一周是1800个点，则划分6段，每段300个点，从中提取20个角点
                        if (largestPickedNum <= 20) {
                            p_cloud_label_[ind] = 1;  // 标记为角点,加入角点点云
                            p_corner_cloud_->push_back(p_extracted_cloud_->points[ind]);
                        } else {
                            break;
                        }

                        p_cloud_neighbor_picked_[ind] = 1;  // 标记已被处理

                        // 同一条扫描线上后5个点标记一下，不再处理，避免特征聚集
                        for (int l = 1; l <= 5; l++) {
                            int columnDiff = std::abs(
                                int(cloud_info_.point_col_ind[ind + l] - cloud_info_.point_col_ind[ind + l - 1]));
                            if (columnDiff > 10)  // 大于10，说明距离远，则不作标记
                                break;
                            p_cloud_neighbor_picked_[ind + l] = 1;
                        }

                        // 同一条扫描线上前5个点标记一下，不再处理，避免特征聚集
                        for (int l = -1; l >= -5; l--) {
                            int columnDiff = std::abs(
                                int(cloud_info_.point_col_ind[ind + l] - cloud_info_.point_col_ind[ind + l + 1]));
                            if (columnDiff > 10)  // 大于10，说明距离远，则不作标记
                                break;
                            p_cloud_neighbor_picked_[ind + l] = 1;
                        }
                    }
                }

                // 按照曲率从小到大遍历
                for (int k = sp; k <= ep; k++) {
                    int ind = cloud_smoothness_vec_[k].ind;  // 激光点索引

                    // 当前激光点还未被处理，且曲率小于阈值，则认为是平面点
                    if (p_cloud_neighbor_picked_[ind] == 0 && p_cloud_curvature_[ind] < surf_threshold_) {
                        p_cloud_label_[ind]           = -1;  // 标记为平面点
                        p_cloud_neighbor_picked_[ind] = 1;   // 标记已被处理

                        // 同一条扫描线上后5个点标记一下，不再处理，避免特征聚集
                        for (int l = 1; l <= 5; l++) {
                            int columnDiff = std::abs(
                                int(cloud_info_.point_col_ind[ind + l] - cloud_info_.point_col_ind[ind + l - 1]));
                            if (columnDiff > 10)  // 大于10，说明距离远，则不作标记
                                break;

                            p_cloud_neighbor_picked_[ind + l] = 1;
                        }

                        // 同一条扫描线上后5个点标记一下，不再处理，避免特征聚集
                        for (int l = -1; l >= -5; l--) {
                            int columnDiff = std::abs(
                                int(cloud_info_.point_col_ind[ind + l] - cloud_info_.point_col_ind[ind + l + 1]));
                            if (columnDiff > 10)  // 大于10，说明距离远，则不作标记
                                break;

                            p_cloud_neighbor_picked_[ind + l] = 1;
                        }
                    }
                }

                // 平面点和未被处理的点(<=0)，都认为是平面点，加入平面点云集合
                for (int k = sp; k <= ep; k++) {
                    if (p_cloud_label_[k] <= 0) {
                        surfaceCloudScan->push_back(p_extracted_cloud_->points[k]);
                    }
                }
            }

            // 平面点云降采样
            surfaceCloudScanDS->clear();
            down_size_filter_.setInputCloud(surfaceCloudScan);
            down_size_filter_.filter(*surfaceCloudScanDS);

            *p_surface_cloud_ += *surfaceCloudScanDS;  // 加入平面点云集合
            // 用surfaceCloudScan来装数据，然后放到downSizeFilter里，
            // 再用downSizeFilter进行.filter()操作，把结果输出到*surfaceCloudScanDS里。
            // 最后把DS装到surfaceCloud中。DS指的是DownSample。
            // 同样角点（边缘点）则没有这样的操作，直接就用cornerCloud来装点云。
        }
    }

    /**
     * 清空准备发布的点云信息中无效的字段，节省数据量
     * 这里清空的数据字段中point_col_ind和point_range的长度都是N_SCAN * N_COL，且下游任务均不使用
     */
    void freeCloudInfoMemory() {
        cloud_info_.start_ring_index.clear();
        cloud_info_.end_ring_index.clear();
        cloud_info_.point_col_ind.clear();
        cloud_info_.point_range.clear();
    }

    /**
     * 发布角点、面点点云，发布带特征点云数据的当前激光帧点云信息
     */
    void publishFeatureCloud() {
        // free cloud info memory
        freeCloudInfoMemory();
        // save newly extracted features
        // 发布角点、面点点云，用于rviz展示
        cloud_info_.cloud_corner =
            publishCloud(p_pub_corner_points_, p_corner_cloud_, cloud_header_.stamp, lidar_frame_);
        cloud_info_.cloud_surface =
            publishCloud(p_pub_surface_points_, p_surface_cloud_, cloud_header_.stamp, lidar_frame_);
        // publish to mapOptimization
        // 发布当前激光帧点云信息，加入了角点、面点点云数据，发布给mapOptimization
        // 和imageProjection.cpp发布的不是同一个话题，
        // image发布的是"lio_sam/deskew/cloud_info"，
        // 这里发布的是"lio_sam/feature/cloud_info"，
        // 因此不用担心地图优化部分的冲突
        p_pub_laser_cloud_info_->publish(cloud_info_);
    }
};

int main(int argc, char **argv) {
    rclcpp::init(argc, argv);
    rclcpp::NodeOptions options;
    options.use_intra_process_comms(true);
    rclcpp::executors::SingleThreadedExecutor exec;

    auto FE = std::make_shared<FeatureExtraction>(options);

    exec.add_node(FE);

    RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "\033[1;32m----> Feature Extraction Started.\033[0m");

    exec.spin();

    rclcpp::shutdown();
    return 0;
}
